from bert4keras.layers import ConditionalRandomField
from keras.layers import LSTM, Masking, Embedding, Input, TimeDistributed, Dropout, Bidirectional
from keras.models import Model
from bert4keras.optimizers import Adam
from crf_layer import CRF



class BilstmCrfModel():

    def __init__(self, max_len, vocab_size, embedding_dim, lstm_units, class_nums):
        super(BilstmCrfModel, self).__init__()
        self.max_len = max_len
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.lstm_units = lstm_units
        self.class_nums = class_nums

    # 创建模型
    def buid_model(self):
        inputs = Input(shape=(self.max_len,))

        x = Masking(mask_value=0)(inputs)

        x = Embedding(input_dim=self.vocab_size, output_dim=self.embedding_dim, trainable=False, mask_zero=True)(x)

        x = Bidirectional(LSTM(units=self.lstm_units, return_sequences=True, kernel_initializer='he_normal'))(x)

        x = TimeDistributed(Dropout(rate=0.2))(x)

        crf = CRF(self.class_nums)
        output = crf(x)

        model = Model(inputs, output)

        model.summary()

        model.compile(
            loss=crf.loss_function,
            optimizer=Adam(5e-5),
            metrics=[crf.accuracy]
        )
        return model
